import torch
import os
import time
import json
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import torchvision.models as models
from matplotlib import pyplot as plt

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ASSETS_DIR = os.path.join(BASE_DIR, r'..\..\..', 'Assets')


def img_transform(img_rgb, transform=None):
    if transform is None:
        raise ValueError('no transform')

    img_t = transform(img_rgb)
    return img_t


def load_class_names(p_clsnames, p_clsname_cn):
    with open(p_clsnames, 'r') as f:
        class_names = json.load(f)
    with open(p_clsname_cn, encoding='UTF-8') as f:
        class_names_cn = f.readlines()
    return class_names, class_names_cn


def get_model(path_state_dict, vis_model=False):
    model = models.alexnet()
    pretrained_state_dict = torch.load(path_state_dict)
    model.load_state_dict(pretrained_state_dict)
    model.eval()

    if vis_model:
        from torchsummary import summary
        summary(model, input_size=(3, 224, 224), device='cpu')

    model.to(device)
    return model


def process_img(path_img):
    norm_mean = [0.485, 0.456, 0.406]
    norm_std = [0.229, 0.224, 0.225]
    inference_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(norm_mean, norm_std),
    ])

    img_rgb = Image.open(path_img).convert('RGB')

    img_tensor = img_transform(img_rgb, inference_transform)
    img_tensor.unsqueeze_(0)
    img_tensor = img_tensor.to(device)

    return img_tensor, img_rgb


if __name__ == '__main__':
    path_state_dict = os.path.join(ASSETS_DIR, 'pre_train_modules', 'pre_train_alexnet.pth')
    path_img = os.path.join(BASE_DIR, '..', 'data', 'tiger cat.jpg')
    path_classnames = os.path.join(BASE_DIR, '..', 'data', 'imagenet1000.json')
    path_classnames_cn = os.path.join(BASE_DIR, '..', 'data', 'imagenet_classnames.txt')

    cls_n, cls_n_cn = load_class_names(path_classnames, path_classnames_cn)

    img_tensor, img_rgb = process_img(path_img)

    alexnet_model = get_model(path_state_dict, True)

    with torch.no_grad():
        time_tic = time.time()
        outputs = alexnet_model(img_tensor)
        time_toc = time.time()

    _,pred_int = torch.max(outputs.data,1)
    _,top5_idx = torch.topk(outputs.data,5,dim=1)

    pred_idx = int(pred_int.cpu().numpy())
    pred_str, pred_cn = cls_n[pred_idx], cls_n_cn[pred_idx]
    print("img: {} is: {}\n{}".format(os.path.basename(path_img),pred_str,pred_cn))
    print("time consuming:{:.2f}s".format(time_toc-time_tic))

    plt.imshow(img_rgb)
    plt.title("predict:{}".format(pred_str))
    top5_num = top5_idx.cpu().numpy().squeeze()
    text_str = [cls_n[t] for t in top5_num]
    for idx in range(len(top5_num)):
        plt.text(5,15+idx*30,'top {}:{}'.format(idx+1,text_str[idx]),bbox=dict(fc='yellow'))
    plt.show()

